
#include "net_asio.h"

using namespace WtfEngine;
using namespace boost::asio::ip;

#ifdef __NET_ASIO__



/**** Generic Socket ****/

Socket::Socket():
	mStatus(UNINITIALIZED), mvBuffer(NULL) {
	NetworkManager::GetInstance().Log(NetworkManager::LOG_DEBUG, "Socket created");
};
Socket::~Socket() {
	NetworkManager::GetInstance().Log(NetworkManager::LOG_DEBUG, "Socket destroyed");
};

void Socket::Listen() {
	switch(mStatus) {
		case CLOSED:
		case UNINITIALIZED:
			// Listen for connections
			NetworkManager::GetInstance().AddSocket(this);
			SetStatus(LISTENING);
			break;
	
		case CONNECTED:
		case RECEIVING:
			// Listen for data
			NetworkManager::GetInstance().AddSocket(this);
			SetStatus(RECEIVING);
			break;

		case SENDING:
			SetStatus(SEND_AND_RECV);
			break;

		default:
			throw new std::logic_error("Network: Cannot listen - socket is not ready");
			break;
	};
};

void Socket::Connect(tIpAddress ip) {
	switch(mStatus) {
		case CLOSED:
		case UNINITIALIZED:
			NetworkManager::GetInstance().AddSocket(this);
			SetStatus(CONNECTING);
			break;
	
		case LISTENING:
			throw new std::logic_error("Cannot connect server socket");
			break;
		case CONNECTED:
			throw new std::logic_error("Socket already connected");
			break;
		default:
			throw new std::logic_error("Network: Cannot connect - socket is not ready");
			break;
	};
};

void Socket::Disconnect() {
	SetStatus(CLOSING);
	NetworkManager::GetInstance().RemoveSocket(this);
};

IPacket::tRef Socket::RecvPacket() {
	IPacket::tRef r;
	
	if(mvRecvQueue.size() > 0) {
		r = mvRecvQueue[0];
		mvRecvQueue.pop_front();
	};
	
	return r;
};

void Socket::OnConnect(const boost::system::error_code &error) {
	if(!error) {
		SetStatus(CONNECTED);
	} else {
		NetworkManager::GetInstance().Log(NetworkManager::LOG_WARNING,
			String("Connection failed - ") + error.message());
		SetStatus(CONNECT_ERROR);
	};
	NetworkManager::GetInstance().RemoveSocket(this);
};

void Socket::OnReceive(const boost::system::error_code& error, std::size_t n) {
	if(!error) {
		if(mStatus == SEND_AND_RECV) {
			mStatus = SENDING;
		} else {
			mStatus = CONNECTED;
		};

		Kernel::GetInstance().ProcessEvent(new SocketDataEvent(this));
	} else {
		NetworkManager::GetInstance().Log(NetworkManager::LOG_WARNING,
			String("Receive failed - ") + error.message());
		SetStatus(RECV_ERROR);
	};

	if(mStatus != SENDING)
		NetworkManager::GetInstance().RemoveSocket(this);
};

void Socket::OnSend(const boost::system::error_code& error) {
	if(!error) {
		if(mStatus == SEND_AND_RECV) {
			SetStatus(RECEIVING);
		} else {
			SetStatus(CONNECTED);
		};
	} else {
		NetworkManager::GetInstance().Log(NetworkManager::LOG_WARNING,
			String("Send failed - ") + error.message());
		SetStatus(SEND_ERROR);
	};
};

void Socket::SetStatus(tStatus st) {
	if(mStatus != st) {
		Kernel::GetInstance().ProcessEvent(new SocketStatusChangeTo(this, st));
	};
	mStatus = st;
};

/**
 * Disconnects the socket and clears all bound event handlers.
 **/
void Socket::Destroy() {
	Disconnect();
	Kernel::GetInstance().GetGameState()->ClearEvents(new SocketEvent(this));
};


/**** TCP ****/

TcpSocket::Server::Server(tPort port)
: Acceptor(NetworkManager::GetInstance().getIoService(), tcp::endpoint(tcp::v4(), port)) {
};

TcpSocket::Server::~Server() {
	Acceptor.close();
};

void TcpSocket::Server::Listen(TcpSocket * pSocket) {
	
	//
	// Create a new TCP socket object to represent the client's connection
	//

	rClientSocket = new TcpSocket(pSocket->getPort());
	rClientSocket->mpClient = new Client();
	rClientSocket->mStatus = ISocket::CONNECTING;
	Acceptor.async_accept(rClientSocket->mpClient->Socket, rClientSocket->mpClient->Endpoint,
		boost::bind(&TcpSocket::OnClientConnect, pSocket, boost::asio::placeholders::error));
};

TcpSocket::Client::Client()
: Socket(NetworkManager::GetInstance().getIoService()), Endpoint() {
};

TcpSocket::Client::~Client() {
	Socket.close();
};

void TcpSocket::Client::Listen(TcpSocket * pSocket) {
	Socket.async_receive(boost::asio::buffer(pSocket->mvBuffer, MaxPacketSize), boost::bind(&TcpSocket::OnReceive, pSocket,
		boost::asio::placeholders::error, boost::asio::placeholders::bytes_transferred));
};

void TcpSocket::Client::Send(WtfEngine::TcpSocket *pSocket, const WtfEngine::IPacket::tRef &rPacket) {
#ifndef NDEBUG
	std::stringstream ss; ss << rPacket->getSize();
	NetworkManager::GetInstance().Log(NetworkManager::LOG_DEBUG, String("TCP Client - sending packet, size ") + ss.str());
#endif

	Socket.async_send(boost::asio::buffer((void *)rPacket->getData(), rPacket->getSize()),
		boost::bind(&TcpSocket::OnSend, pSocket, boost::asio::placeholders::error));
};


TcpSocket::TcpSocket(WtfEngine::tPort port)
: Socket(), mpClient(NULL), mPort(port) {
};

TcpSocket::~TcpSocket() {
	if(mStatus == LISTENING) {
		delete mpServer;
	} else if(mpClient) {
		delete mpClient;
	};

	if(mvBuffer) delete [] mvBuffer;
};

void TcpSocket::Listen() {
	Socket::Listen();
	if(mStatus == LISTENING) {
		// Server socket - listen for connections
		assert(mpServer == NULL);
		mpServer = new Server(mPort);
		mpServer->Listen(this);
	} else {
		// Client socket - listen for data
		assert(mpClient && !mvBuffer);
		mvBuffer = new unsigned char[MaxPacketSize];
		NetworkManager::GetInstance().AddSocket(this);
		mpClient->Listen(this);
	};
};

void TcpSocket::Connect(tIpAddress ip) {
	Socket::Connect(ip);
	assert(mpClient == NULL);
	mpClient = new Client();

	mpClient->Socket.async_connect(mpClient->Endpoint, boost::bind(&TcpSocket::OnConnect, this,
		boost::asio::placeholders::error));
};

void TcpSocket::Disconnect() {
	if(mStatus == LISTENING) {
		delete mpServer;
	} else if(mpClient) {
		delete mpClient;
	};

	if(mvBuffer) {
		delete [] mvBuffer;
		mvBuffer = NULL;
	};

	Socket::Disconnect();
	SetStatus(CLOSED);
};

void TcpSocket::SendPacket(const IPacket::tRef& rPacket) {
	mvSendQueue.push_back(rPacket);
	if(mStatus == CONNECTED) {
		// Add to network manager to keep alive whilst sending.
		NetworkManager::GetInstance().AddSocket(this);
		mStatus = SENDING;
		mpClient->Send(this, rPacket);
	} else if(mStatus != SENDING) {
		throw new std::logic_error("Network: cannot send packet - socket not ready");
	};
};

void TcpSocket::OnReceive(const boost::system::error_code& error, std::size_t n) {
	// Create packet
	mvRecvQueue.push_back(new NonCopyPacket(mvBuffer, n));
	mvBuffer = NULL;

	Socket::OnReceive(error, n);
};

void TcpSocket::OnConnect(const boost::system::error_code& error) {
	Socket::OnConnect(error);
};

void TcpSocket::OnClientConnect(const boost::system::error_code& error) {
	if(!error) {
		mpServer->rClientSocket->OnConnect(error);
		Kernel::GetInstance().ProcessEvent(new SocketClientConnectEvent(this, mpServer->rClientSocket));

		// Rebind
		if(mStatus == LISTENING) {
			mpServer->Listen(this);
		} else {
			NetworkManager::GetInstance().RemoveSocket(this);
		};
	};
};

void TcpSocket::OnSend(const boost::system::error_code& error) {
	// Pop packet from queue, send next if there is one
	mvSendQueue.pop_front();
	Socket::OnSend(error);
	
	if(mStatus == CONNECTED && !mvSendQueue.empty()) {
		mStatus = SENDING;
		mpClient->Send(this, mvSendQueue[0]);
	} else if(mStatus == RECEIVING && !mvSendQueue.empty()) {
		mStatus = SEND_AND_RECV;
		mpClient->Send(this, mvSendQueue[0]);
	} else {
		NetworkManager::GetInstance().RemoveSocket(this);
	};
};


/**** UDP ****/

UdpSocket::UdpSocket(WtfEngine::tPort port)
: Socket(), mSocket(NetworkManager::GetInstance().getIoService()), mEndpoint(udp::v4(), port) {
};

UdpSocket::~UdpSocket() {
	if(mvBuffer) delete mvBuffer;

	mSocket.close();
};


/**
 * Although UDP is connectionless, this initializes the socket to send data to the specified
 * address.
 **/
void UdpSocket::Connect(tIpAddress ip) {
	Socket::Connect(ip);

	mEndpoint = udp::endpoint(address_v4(ip), mEndpoint.port());
	mSocket.connect(mEndpoint);

	SetStatus(CONNECTED);
};

void UdpSocket::Listen() {
	// Can not listen on "connected" UDP sockets.
	assert(mStatus != CONNECTED);
	Socket::Listen();

	assert(!mvBuffer);
	mvBuffer = new unsigned char[MaxPacketSize];
	NetworkManager::GetInstance().AddSocket(this);
	mSocket.async_receive(boost::asio::buffer(mvBuffer, MaxPacketSize), boost::bind(&UdpSocket::OnReceive, this,
		boost::asio::placeholders::error, boost::asio::placeholders::bytes_transferred));
};

void UdpSocket::Disconnect() {
	Socket::Disconnect();
	mSocket.close();

	if(mvBuffer) {
		delete [] mvBuffer;
		mvBuffer = NULL;
	};

	SetStatus(CLOSED);
};

void UdpSocket::SendPacket(const WtfEngine::IPacket::tRef &rPacket) {
	mvSendQueue.push_back(rPacket);
	if(mStatus == CONNECTED) {
		// Add to network manager to keep alive whilst sending.
		NetworkManager::GetInstance().AddSocket(this);
		mStatus = SENDING;

		#ifndef NDEBUG
			std::stringstream ss; ss << rPacket->getSize();
			NetworkManager::GetInstance().Log(NetworkManager::LOG_DEBUG, String("UDP Broadcast - sending packet, size ") + ss.str());
		#endif

		mSocket.async_send(boost::asio::buffer((void *)rPacket->getData(), rPacket->getSize()),
			boost::bind(&UdpSocket::OnSend, this, boost::asio::placeholders::error));
	} else if(mStatus != SENDING) {
		throw new std::logic_error("Network: cannot send packet - socket not ready");
	};
};

void UdpSocket::OnReceive(const boost::system::error_code& error, std::size_t n) {
	if(!error) {
		// Create packet
		mvRecvQueue.push_back(new NonCopyPacket(mvBuffer, n));
	};

	mvBuffer = NULL;
	Socket::OnReceive(error, n);
};

void UdpSocket::OnSend(const boost::system::error_code &error) {
	// Pop packet from queue, send next if there is one
	mvSendQueue.pop_front();
	Socket::OnSend(error);
	
	if(mStatus == CONNECTED && !mvSendQueue.empty()) {
		mStatus = SENDING;

		#ifndef NDEBUG
			std::stringstream ss; ss << mvSendQueue[0]->getSize();
			NetworkManager::GetInstance().Log(NetworkManager::LOG_DEBUG, String("UDP Broadcast - sending packet, size ") + ss.str());
		#endif

		mSocket.async_send(boost::asio::buffer((void *)mvSendQueue[0]->getData(), mvSendQueue[0]->getSize()),
			boost::bind(&UdpSocket::OnSend, this, boost::asio::placeholders::error));
	} else {
		NetworkManager::GetInstance().RemoveSocket(this);
	};
};



/**** Broadcast ****/

UdpBroadcastSocket::UdpBroadcastSocket(tPort port)
: UdpSocket(port)
{
	UdpSocket::Connect(tIpAddress(address_v4::broadcast().to_ulong()));
};

/**** Network Manager ****/

NetworkManager::NetworkManager():
		Task<NetworkManager>(0, "Network"), Loggable<NetworkManager>("Network"), mIoService() {
};

void NetworkManager::Run() {
	// Poll the IO Service to dispatch network handlers
	mIoService.poll();
};

#endif
